/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef ATB_PLUGIN_ACLNN_NN_OPERATION_H
#define ATB_PLUGIN_ACLNN_NN_OPERATION_H
#include <string>
#include <acl/acl.h>
#include <aclnn/acl_meta.h>
#include "atb/atb_infer.h"
#include "atb/operation.h"
#include "utils.h"
#include "acl_nn_tensor.h"

namespace atb_plugin {

class AclNNOperation : public atb::Operation {
public:
    explicit AclNNOperation(const std::string &opName);
    ~AclNNOperation() override;
    std::string GetName() const override;
    atb::Status Setup(const atb::VariantPack &variantPack, uint64_t &workspaceSize, atb::Context *context) override;
    atb::Status Execute(const atb::VariantPack &variantPack, uint8_t *workspace, uint64_t workspaceSize,
                        atb::Context *context) override;
    void DestoryOperation();
    atb::Status CreateAclNNOpCache(const atb::VariantPack &variantPack);
    virtual int CreateAclNNVariantPack(const atb::VariantPack &variantPack) = 0;
    virtual int SetAclNNWorkspaceExecutor() = 0;
    virtual int ExecuteAclNNOp(uint8_t *workspace, aclrtStream &stream) = 0;
    std::string opName_;
    atb::SVector<std::shared_ptr<AclNNTensor>> aclInTensors;
    atb::SVector<std::shared_ptr<AclNNTensor>> aclOutTensors;
    aclOpExecutor* aclExecutor;
    uint64_t workspaceSize;
};
} // namespace atb_plugin
#endif